/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "graph_optimizer/op_compiler/op_compiler_optune.h"

namespace fe{

OpCompilerOpTune::OpCompilerOpTune(const std::string& compiler_name, const std::string& engine_name,
                                   OpStoreAdapterManagerPtr op_store_adapter_manager_ptr)
    :OpCompiler(compiler_name, engine_name, op_store_adapter_manager_ptr){}

OpCompilerOpTune::~OpCompilerOpTune() {}

Status OpCompilerOpTune::Initialize(BufferFusionFunc &func) {
  // if graph optimizer has been initialized, return SUCCESS
  if (init_flag_) {
    FE_LOGW("OpCompiler has been initialized.");
    return SUCCESS;
  }

  init_flag_ = true;
  bufferFusionFunc_ = func;
  return SUCCESS;
}

Status OpCompilerOpTune::RunCompileProcess(ge::ComputeGraph& graph, const std::shared_ptr<GraphComm>& graph_comm_ptr,
                                           const std::shared_ptr<ScopeAllocator>& scope_allocator_ptr,
                                           std::vector<ge::NodePtr> &buff_fus_compile_failed_nodes,
                                           bool &need_post_process) {
  /* Some nodes needs to be re-pre-compiled after Ub fusion matching.
   * Because there format or data type. */
  Status ret;
  bool need_re_precompile_graph = false;
  (void)ge::AttrUtils::GetBool(graph, NEED_RE_PRECOMPILE, need_re_precompile_graph);
  if (need_re_precompile_graph) {
    ret = PreCompileOp(graph);
    if (ret != SUCCESS) {
      REPORT_FE_ERROR("[SubGraphOpt][Compile][Optune] PreCompileOp failed after buffer fusion for graph %s.",
                      graph.GetName().c_str());
      return ret;
    }
  }

  CompileInfoParam compile_info(buff_fus_compile_failed_nodes);
  vector<ge::NodePtr> nodes_be_compiled;
  vector<ge::NodePtr> buff_fus_rollback_nodes;

  ret = OpCompiler::GetFusionScope(graph, buff_fus_rollback_nodes, compile_info.fusion_nodes_map, nodes_be_compiled);
  if (ret != SUCCESS) {
    return ret;
  }

  /* This compile is for getting a graph which will be successfully compiled(for those failed fusion-nodes, we
   * will roll it back to single node). In ms-tune process, graph also goes through this process. So for
   * matching the lx-fusion strategy in op-tune, we need to compile and roll-back first. */
  compile_info.ignore_compile_strategy = true;
  ret = CompileOpOnly(graph, compile_info);
  if (ret != SUCCESS) {
    return ret;
  }

  FE_CHECK_NOTNULL(bufferFusionFunc_);

  /* Call lx-fusion to optimize the graph. */
  ret = bufferFusionFunc_(graph, graph_comm_ptr, scope_allocator_ptr);
  if (ret != SUCCESS && ret != tune::NO_FUSION_STRATEGY) {
    REPORT_FE_ERROR("[SubGraphOpt][Compile][Optune] BufferFusionProcess failed for graph %s result %u.",
                    graph.GetName().c_str(), ret);
    return FAILED;
  }

  // In op tuning mode, precompile only needs to be done with all new
  // nodes which are generated by lx-fusion.
  // And the compile is necessary for those which are set attribute
  // "compile_strategy". For now, because the compilation time cost is not
  // that important, we just compile all nodes. And module Te-fusion will differentiate
  // compile_strategy.
  if (ret != tune::NO_FUSION_STRATEGY) {
    ret = PreCompileOp(graph);
    if (ret != SUCCESS) {
      REPORT_FE_ERROR("[SubGraphOpt][Compile][Optune] Re-precompile failed after lx fusion for graph %s.",
                      graph.GetName().c_str());
      return ret;
    }
  }

  compile_info.fusion_nodes_map.clear();
  compile_info.scope_json_map.clear();
  compile_info.buff_fus_to_del_nodes.clear();
  compile_info.ignore_compile_strategy = false;
  nodes_be_compiled.clear();

  ret = OpCompiler::GetFusionScope(graph, buff_fus_rollback_nodes,
                                   compile_info.fusion_nodes_map, nodes_be_compiled);
  if (ret != SUCCESS) {
    return FAILED;
  }
  /* Here we try to re-compile every op to do the op-tune.
   * This step is to generate json for every nodes. */
  ret = CompileOpOnly(graph, compile_info);
  if (ret != SUCCESS) {
    return ret;
  }
  /* In op tune, we do not care the result of compilation. So just return. */
  need_post_process = false;
  return SUCCESS;
}
}